#include "TinyEngine.h"
#include <assert.h>

TinyEngine::TinyEngine(uint32_t k0, uint32_t k1, uint32_t k2, uint32_t k3, size_t rounds = 32)
{
	key.kvar.k0 = k0;
	key.kvar.k1 = k1;
	key.kvar.k2 = k2;
	key.kvar.k3 = k3;
	numberOfRounds = rounds;
}

bool TinyEngine::process(DataSource* source, DataTarget* target, std::string pswd, EncryptionMode mode)
{
	while (source->hasNextByte())
	{
		size_t cnt = getBlock(source);
		assert(block_size == cnt);
		
		// process
		bool res = TinyEngine::processPair(mode);
		assert(res);

		writeBlock(target);
	}
	return true;
}

size_t TinyEngine::getBlock(DataSource *source)
{
	size_t r = 0;
	uint8_t* vv = (uint8_t*)data.v;
	for (size_t i = 0; i < block_size; ++i)
	{
		if (source->hasNextByte())
		{
			vv[i] = source->getNextByte();
			++r;
		}
		else 
			break;
	}
	return r;
}

uint32_t TinyEngine::writeBlock(DataTarget *target)
{
	size_t r = 0;
	uint8_t* vv = (uint8_t*)data.v;
	for (size_t i = 0; i < block_size; ++i)
	{
		r += (int)target->putNextByte(*vv);
		++vv;
	}
	return r;
}

bool TinyEngine::processPair(EncryptionMode mode)
{
	sum = EncryptionMode::Enctypt == mode ? 0 : 0xC6EF3720;
	for (size_t i = 0; i < numberOfRounds; ++i)
	{
		switch (mode)
		{
			case Enctypt:
				encrypt();
				break;
			case Decrypt:
				decrypt();
				break;
			default:
				assert(0);
		}
	}
	return true;
}

void TinyEngine::encrypt()
{
	sum += delta;
	data.vvar.v0 += ((data.vvar.v1 << 4) + key.kvar.k0) ^ (data.vvar.v1 + sum) ^ ((data.vvar.v1 >> 5) + key.kvar.k1);
	data.vvar.v1 += ((data.vvar.v0 << 4) + key.kvar.k2) ^ (data.vvar.v0 + sum) ^ ((data.vvar.v0 >> 5) + key.kvar.k3);
}

void TinyEngine::decrypt()
{
	data.vvar.v1 -= ((data.vvar.v0 << 4) + key.kvar.k2) ^ (data.vvar.v0 + sum) ^ ((data.vvar.v0 >> 5) + key.kvar.k3);
	data.vvar.v0 -= ((data.vvar.v1 << 4) + key.kvar.k0) ^ (data.vvar.v1 + sum) ^ ((data.vvar.v1 >> 5) + key.kvar.k1);
	sum -= delta;
}